import cv2
import matplotlib.pyplot as plt
import numpy as np
import random

def rotate_bound(image, angle):
    # grab the dimensions of the image and then determine the
    # center
    (h, w) = image.shape[:2]
    (cX, cY) = (w // 2, h // 2)
 
    # grab the rotation matrix (applying the negative of the
    # angle to rotate clockwise), then grab the sine and cosine
    # (i.e., the rotation components of the matrix)
    M = cv2.getRotationMatrix2D((cX, cY), -angle, 1.0)
    cos = np.abs(M[0, 0])
    sin = np.abs(M[0, 1])
 
    # compute the new bounding dimensions of the image
    nW = int((h * sin) + (w * cos))
    nH = int((h * cos) + (w * sin))
 
    # adjust the rotation matrix to take into account translation
    M[0, 2] += (nW / 2) - cX
    M[1, 2] += (nH / 2) - cY
 
    # perform the actual rotation and return the image
    return cv2.warpAffine(image, M, (nW, nH), borderValue = (255*random.random(),255*random.random(),255*random.random()))

def mergeImg(inputImg,maskImg,contourData,drawPosition):
    '''
    :param inputImg: 输入的图像
    :param maskImg: 输入的模板图像
    :param contourData: 输入的模板中轮廓数据 numpy 形式如[(x1,y1),(x2,y2),...,]
    :param drawPosition: （x,y） 大图中要绘制模板的位置,以maskImg左上角为起始点
    :return: outPutImg：输出融合后的图像
             outContourData: 输出轮廓在inputImg的坐标数据
             outRectData: 输出轮廓的矩形框在inputImg的坐标数据
    '''
    #通道需要相等
    if (inputImg.shape[2] != maskImg.shape[2]):
        print("inputImg shape != maskImg shape")
        return
    inputImg_h=inputImg.shape[0]
    inputImg_w=inputImg.shape[1]
    maskImg_h = maskImg.shape[0]
    maskImg_w = maskImg.shape[1]
    #inputImg图像尺寸不能小于maskImg
    if(inputImg_h<maskImg_h or inputImg_w<maskImg_w):
        print("inputImg size < maskImg size")
        return
    #画图的位置不能超过原始图像
    if(((drawPosition[0]+maskImg_w)>inputImg_w) or ((drawPosition[1]+maskImg_h)>inputImg_h)):
        print("drawPosition + maskImg > inputImg range")
        return
    outPutImg=inputImg.copy()
    input_roi=outPutImg[drawPosition[1]:drawPosition[1]+maskImg_h,drawPosition[0]:drawPosition[0]+maskImg_w]  #position
    imgMask_array=np.zeros((maskImg_h,maskImg_w,maskImg.shape[2]),dtype=np.uint8)
    
    # triangles_list=[contourData]
    # cv2.fillPoly(imgMask_array, triangles_list, color=(1,1,1))
    # cv2.fillPoly(input_roi, triangles_list, color=(0, 0, 0))
    
    # imgMask_array=imgMask_array*maskImg
    weight_mix = (random.random() + 0.3)*0.5
    output_ori=input_roi*(1-weight_mix)+maskImg*weight_mix
    # output_ori=input_roi+imgMask_array
    outPutImg[drawPosition[1]:drawPosition[1] + maskImg_h, drawPosition[0]:drawPosition[0] + maskImg_w]=output_ori
    # triangles_list[0][:, 0] = contourData[:, 0] +drawPosition[0]
    # triangles_list[0][:, 1] = contourData[:, 1] +drawPosition[1]
    # outContourData=triangles_list[0]
    outContourData = 0
    return outPutImg,outContourData#,outRectData


# import pycocotools.mask as mask_util
# from PIL import Image


cap = cv2.VideoCapture('/media/st/Application/Ubuntu/dataset/dangerous0731/1.mp4')
sign_img1=cv2.imread('/media/st/Application/Ubuntu/dataset/dangerous0731/sign1.jpg')
sign_img1 = cv2.resize(sign_img1, dsize=(sign_img1.shape[1]//20, sign_img1.shape[0]//20)) 
# sign_img2=cv2.imread('/media/st/Application/Ubuntu/dataset/dangerous0731/sign2.jpg')
# sign_img2 = cv2.resize(sign_img2, dsize=(sign_img2.shape[1]//5, sign_img2.shape[0]//5)) 
sign_img3=cv2.imread('/media/st/Application/Ubuntu/dataset/dangerous0731/sign3.jpg')
sign_img3 = cv2.resize(sign_img3, dsize=(sign_img3.shape[1]//5, sign_img3.shape[0]//5)) 
frame_num_ = 0
cap_num_ = 10

while(cap.isOpened()):
    cap_num_ = cap_num_ +1
    ret, frame = cap.read()
    if cap_num_ % 10 != 0:
        continue

    #scale resize
    frame = cv2.resize(frame, dsize=(frame.shape[1]//2, frame.shape[0]//2))  # dsize的输入必须为整型

    for i in range(2):
        frame_temp = frame.copy()
        sign_img1_temp = sign_img1.copy()
        sign_img1_resize = cv2.resize(sign_img1_temp, dsize=( int(sign_img1_temp.shape[1]*(random.random()+0.5)), int(sign_img1_temp.shape[0]*(random.random()+0.5))))

        rotate_angle = int(random.random()*360)
        sign_img1_rotate = rotate_bound(sign_img1_resize, rotate_angle)

        drawPosition =( int((frame.shape[1]-sign_img1_rotate.shape[1]-1)*random.random()), int((frame.shape[0]-sign_img1_rotate.shape[0]-1)*random.random())) # xy  == w h
        # print("debug postion ", drawPosition)
        print("debug num  ", frame_num_)
        cenX = float( (drawPosition[0] + sign_img1_rotate.shape[1] *0.5) / frame_temp.shape[1] )
        cenY = float( (drawPosition[1] + sign_img1_rotate.shape[0] *0.5) / frame_temp.shape[0] )
        bbox_w  = float(sign_img1_rotate.shape[1]) / frame_temp.shape[1]
        bbox_h  = float(sign_img1_rotate.shape[0]) / frame_temp.shape[0]

        file = open('/media/st/Application/Ubuntu/dataset/dangerous0731/label/'+str(frame_num_)+'.txt', 'w') 
        file.write(str(1)+' '+ str(cenX) + ' ' + str(cenY) + ' ' + str(bbox_w) + ' ' + str(bbox_h) +' ' + '\n')
        
        
        contourData=np.array([(57,7),(107,30),(107,120),(62,122),(2,95),(9,32)])
        outframe,_ = mergeImg(frame_temp, sign_img1_rotate, contourData ,drawPosition)

        cv2.imwrite('/media/st/Application/Ubuntu/dataset/dangerous0731/img/'+str(frame_num_)+'.jpg', outframe)
        frame_num_=frame_num_ + 1

    for i in range(2):
        frame_temp = frame.copy()
        sign_img3_temp = sign_img3.copy()
        sign_img3_resize = cv2.resize(sign_img3_temp, dsize=( int(sign_img3_temp.shape[1]*(random.random()+0.5)), int(sign_img3_temp.shape[0]*(random.random()+0.5))))

        rotate_angle = int(random.random()*360)
        sign_img3_rotate = rotate_bound(sign_img3_resize, rotate_angle)

        drawPosition =( int((frame.shape[1]-sign_img3_rotate.shape[1]-1)*random.random()), int((frame.shape[0]-sign_img3_rotate.shape[0]-1)*random.random())) # xy  == w h
        # print("debug postion ", drawPosition)
        print("debug num  ", frame_num_)
        cenX = float( (drawPosition[0] + sign_img3_rotate.shape[1] *0.5) / frame_temp.shape[1] )
        cenY = float( (drawPosition[1] + sign_img3_rotate.shape[0] *0.5) / frame_temp.shape[0] )
        bbox_w  = float(sign_img3_rotate.shape[1]) / frame_temp.shape[1]
        bbox_h  = float(sign_img3_rotate.shape[0]) / frame_temp.shape[0]

        file = open('/media/st/Application/Ubuntu/dataset/dangerous0731/label/'+str(frame_num_)+'.txt', 'w') 
        file.write(str(1)+' '+ str(cenX) + ' ' + str(cenY) + ' ' + str(bbox_w) + ' ' + str(bbox_h) +' ' + '\n')
        
        
        contourData=np.array([(57,7),(107,30),(107,120),(62,122),(2,95),(9,32)])
        outframe,_ = mergeImg(frame_temp, sign_img3_rotate, contourData ,drawPosition)

        cv2.imwrite('/media/st/Application/Ubuntu/dataset/dangerous0731/img/'+str(frame_num_)+'.jpg', outframe)
        frame_num_=frame_num_ + 1

    


    # cv2.imshow('frame',gray)
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break




cap.release()
cv2.destroyAllWindows()


#scale resize
# image_resize = cv2.resize(image, dsize=(w//2, h//2))  # dsize的输入必须为整型